import os
os.environ["KERAS_BACKEND"] = "torch"
import numpy as np
import pandas as pd
from keras.models import Model
from keras.layers import Embedding, Dense, Input, GRU, AdditiveAttention
# from tensorflow.python.keras.model import Model
# from tensorflow.python.keras.layers import Embedding, Dense, Input, GRU, AdditiveAttention
from text_cnn_test import get_text_labels, PaddingTokenizer, train_test_split, np_utils, EarlyStopping
from _attention import Attention


def build_model(embed_dim: int, vocab_size: int, num_classes: int, activation='softmax'):
	inputs = Input(shape=(None,))
	out = Embedding(vocab_size, embed_dim, input_length=None, mask_zero=False)(inputs)
	out = GRU(units=128, return_sequences=True)(out)
	out = Attention(name="attention")(out)
	out = Dense(num_classes, activation=activation)(out)
	return Model(inputs=inputs, outputs=out)


if __name__ == '__main__':
	file = '~/project/python/parttime/text_gcn/data/北方地区不安全事件统计20240331.csv'
	data = pd.read_csv(file, encoding='GBK')
	texts, labels, classes = get_text_labels(data, text_col='故障描述', label_col='故障标志')
	del data
	
	stop_words = [str(i) for i in range(1)]
	stop_words.extend(['-', '的', '一', '是', '了', '年', '月', '日'])
	tokenizer = PaddingTokenizer(texts=texts, min_freq=5, cut_type='char', stop_words=stop_words, word_freq=False)
	ids = tokenizer.batch_encode(texts)
	X = np.array(ids)
	y = np.array(labels)  # [int]
	# 以下功能相同，也可以不加
	# y = np.expand_dims(y, axis=1)
	# y = y.reshape(-1, 1)
	
	y = np_utils.to_categorical(y, len(classes))  # loss用categorical_crossentropy

	X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
	model = build_model(128, tokenizer.vocab_size, len(classes))
	# print(model.summary())
	early_stopping = EarlyStopping(monitor='val_loss', patience=3)
	model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
	# model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
	model.fit(X_train, y_train, batch_size=32, epochs=50, validation_data=(X_test, y_test), callbacks=early_stopping)